import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules

# 加载数据集
try:
    df = pd.read_excel('餐厅数据.xlsx')
except FileNotFoundError:
    print("找不到指定的文件，请检查路径是否正确。")
    exit(1)

# 转换菜品数据格式为事务列表
transactions = df['菜品'].str.split(',').to_list()

# 使用TransactionEncoder转换事务列表
te = TransactionEncoder()
te_ary = te.fit(transactions).transform(transactions)
df_encoded = pd.DataFrame(te_ary, columns=te.columns_)

# 使用Apriori算法生成频繁项集
frequent_itemsets = apriori(df_encoded, min_support=0.1, use_colnames=True)
frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)

# 生成关联规则
rules = association_rules(frequent_itemsets, metric='confidence', min_threshold=0.1)

# 筛选有效规则：lift > 1 和 confidence > 0.1
effective_rules = rules[(rules['lift'] > 1) & (rules['confidence'] > 0.1)].sort_values(by=['lift', 'confidence'], ascending=False)

# 将频繁项集和筛选后的规则保存到文件
try:
    frequent_itemsets.to_pickle("frequent_itemsets.pkl")
    effective_rules.to_pickle("rule.pkl")
    print("频繁项集和规则已成功保存。")
except Exception as e:
    print(f"保存过程中发生错误: {e}")

# 注意：请根据实际情况调整文件路径以及支持度(min_support)、置信度(min_threshold)阈值。